#include "ImageSet.h"

#include <iostream>
using namespace std;

//////////////////////////////////////////////////////////////////
// ImageSet Object												//
//////////////////////////////////////////////////////////////////

// Constructor for a set of images. This constructor creates an 
// object which contains all the images from a specified folder.
ImageSet::ImageSet(char* directory)
{
	_n = 0;
	_head = NULL;
	_tail = NULL;
	AddToSet(directory,"*.png");
	AddToSet(directory,"*.jpg");
	AddToSet(directory,"*.jpeg");
	AddToSet(directory,"*.bmp");
}

// Constructor for a set of images. This constructor creates a set
// of images from a a single image object. This is useful when a
// set must be constructed from images in memory instead of images
// stored in a folder.
ImageSet::ImageSet(Image* img)
{
	_n = 1;
	_head = new ImageNode;
	_head->_image = new Image(*img);
	_head->next = NULL;
	_head->previous = NULL;
	_tail = _head;
	_homogeneous = true;
}

// Copy constructor for set of images. This copy constructor is 
// used to create a new set of images from an existing set. Note 
// that this is not trivial due to the use of dynamic memory. 
ImageSet::ImageSet(const ImageSet &init)
{
	// Copy across number of images and homogeneity indicator
	_n = init._n;
	_homogeneous = init._homogeneous;

	// Check if we need to copy stuff across
	if(init._head)
	{
		// Initialise the beginning of the list with the correct image
		_head = new ImageNode;
		if(init._head->_image)
			_head->_image = new Image(*(init._head->_image));
		_head->next = NULL;
		_head->previous = NULL;
		_tail = _head;


		ImageNode* temp = init._head;

		// Traverse the list and copy everything across
		while(temp->next)
		{
			_tail->next = new ImageNode;
			_tail->next->_image = new Image(*(temp->next->_image));
			_tail->next->next = NULL;
			_tail->next->previous = _tail;
			_tail = _tail->next;
			temp = temp->next;
		}
	}
	else
	{
		_head = NULL;
		_tail = NULL;
	}
}

// Destructor for set of images. This destructor clears up all 
// memory allocated for images when an object is destroyed.
ImageSet::~ImageSet()
{
	// Loop while we are not at the head node
	if(_tail)
		while(_tail->previous)
		{
			// Jump back one
			_tail = _tail->previous;
			// Delete everthing from the node at the end of the list
			if(_tail->next->_image)
				delete _tail->next->_image;
			if(_tail->next)
				delete _tail->next;
		}

	// Delete the head node
	if(_head)
	{
		if(_head->_image)
			delete _head->_image;
		delete _head;
	}
}

// This function adds files to a set of images from a specified 
// folder with a specified file extension. It returns a boolean 
// value which is true if files were added, and false if none were 
// found.
bool ImageSet::AddToSet(char* directory, char* ext)
{
	// Create string for finding image files
	char file_find[100];
	strcpy_s(file_find,directory);
	strcat_s(file_find,ext);

	// Create string for file path name
	char current_file[100];

	// Get data for first file
	WIN32_FIND_DATA findFileData;
	HANDLE hFind = FindFirstFile(file_find, &findFileData);
	if(hFind == INVALID_HANDLE_VALUE)
		return false;
	strcpy_s(current_file,directory);
	strcat_s(current_file,findFileData.cFileName);

	// Check if any files have already been added
	if(_head)
	{
		// Add to the exisiting list
		_tail->next = new ImageNode;
		_tail->next->_image = new Image;
		_tail->next->_image->Load(current_file);
		_tail->next->next = NULL;
		_tail->next->previous = _tail;
		_tail = _tail->next;
		_n++;
		if(_homogeneous)
		{
			if(	(_head->_image->GetBPP() != _tail->_image->GetBPP()) ||
				(_head->_image->GetWidth() != _tail->_image->GetWidth()) ||
				(_head->_image->GetHeight() != _tail->_image->GetHeight())	)
				_homogeneous = false;
		}
	}
	else
	{
		// Create head of list
		_head = new ImageNode;
		_head->_image = new Image;
		_head->_image->Load(current_file);
		_head->next = NULL;
		_head->previous = NULL;
		_tail = _head;
		_n = 1;
		_homogeneous = true;
	}

	// Fill the rest of the list with the remaining files
	while(FindNextFile(hFind,&findFileData))
	{
		strcpy_s(current_file,directory);
		strcat_s(current_file,findFileData.cFileName);
		_tail->next = new ImageNode;
		_tail->next->_image = new Image;
		_tail->next->_image->Load(current_file);
		_tail->next->next = NULL;
		_tail->next->previous = _tail;
		_tail = _tail->next;
		_n++;
		if(_homogeneous)
		{
			if(	(_head->_image->GetBPP() != _tail->_image->GetBPP()) ||
				(_head->_image->GetWidth() != _tail->_image->GetWidth()) ||
				(_head->_image->GetHeight() != _tail->_image->GetHeight())	)
				_homogeneous = false;
		}
	}
	return true;
}

// This function adds an image from memory to an image set.
void ImageSet::AddToSet(Image* img)
{
	_tail->next = new ImageNode;
	_tail->next->_image = new Image(*img);
	_tail->next->next = NULL;
	_tail->next->previous = _tail;
	_tail = _tail->next;
	_n++;
	if(_homogeneous)
	{
		if(	(_head->_image->GetBPP() != _tail->_image->GetBPP()) ||
			(_head->_image->GetWidth() != _tail->_image->GetWidth()) ||
			(_head->_image->GetHeight() != _tail->_image->GetHeight())	)
			_homogeneous = false;
	}
}

// This function saves a set of images to the specified directory.
void ImageSet::SaveSet(char* directory)
{
	if(_head)
	{
		char buffer[10];
		char file_name[20];
		char file_save[100];
		cout << "Saving images..." << endl;
		ImageNode* temp = _head;
		int i = 1;
		while(temp)
		{
			_itoa(i,buffer,10);
			strcpy_s(file_name,buffer);
			strcpy_s(file_save,directory);
			strcat_s(file_name,".png");
			strcat_s(file_save,file_name);
			temp->_image->Save(file_save,Gdiplus::ImageFormatPNG);
			temp = temp->next;
			i++;
		}
	}
	else
		cout << "Cannot save images, as no images are loaded." << endl;
}

// This function overwrites an image set with greyscale versions
// of those images.
void ImageSet::ConvertSetToGreyscale()
{
	if(_head)
	{
		// Create some variables to check if homongeneity still holds.
		int bpp, w, h;
		bool temp_homogeneous = true;
		cout << "Converting to greyscale..." << endl;
		ImageNode* temp = _head;
		while(temp)
		{
			temp->_image->ConvertToGreyscale();
			// Homogeneity must be checked if false, as it could 
			// change from false to true if the images differ only 
			// in their colour depth.
			if(!_homogeneous)
			{
				if(temp == _head)
				{
					bpp = temp->_image->GetBPP();
					w = temp->_image->GetWidth();
					h = temp->_image->GetHeight();
				}
				else
				{
					if(temp_homogeneous)
					{
						if(	(bpp != temp->_image->GetBPP()) ||
							(w != temp->_image->GetWidth()) ||
							(h != temp->_image->GetHeight())	)
							temp_homogeneous = false;
					}
				}
			}
			temp = temp->next;
		}
		cout << "Conversion complete" << endl;
	}
	else
		cout << "Cannot run conversion, as no images are loaded." << endl;
}

// This function returns the average image of a set. Note that all 
// the images must have the same dimensions.
Image* ImageSet::GetAverageImage()
{
	if(_head)
	{
		// Find the average image
		if(_homogeneous)
		{
			int bpp = _head->_image->GetBPP();
			int w = _head->_image->GetWidth();
			int h = _head->_image->GetHeight();

			Image* avg_img = new Image;
			avg_img->Create(w,h,bpp);
			int* avg = new int[3*w*h];

			// Initialise image pixels to zero
			int* avg_p = avg;
			for(int i=0; i<3*w*h; i++)
				*avg_p++ = 0;

			// Iterate through images in the set
			ImageNode* temp = _head;
			while(temp)
			{
				avg_p = avg;				// Reset the pointer to the start
				BYTE* pixel = (BYTE*)temp->_image->GetPixelAddress(0,temp->_image->GetHeight()-1);
				// The pitch is not the same as the image width, therefore we 
				// must check to see if any extra bytes are stored beyond the
				// end of the image row.
				int buffer = abs(temp->_image->GetPitch()) % (temp->_image->GetBPP()/8);
				for(int i=0; i<temp->_image->GetHeight(); i++)
				{
					for(int j=0; j<temp->_image->GetWidth(); j++)
					{
						*avg_p++ = *avg_p + (int)*pixel++;
						*avg_p++ = *avg_p + (int)*pixel++;
						*avg_p++ = *avg_p + (int)*pixel++;
					}
					for(int k=0; k<buffer; k++)
					{
						pixel++;
					}
				}
				temp = temp->next;
			}

			// Copy into another class
			avg_p = avg;					// Reset the pointer to the start
			for(int i=0; i<3*w*h; i++)
				*avg_p++ = *avg_p/_n;
			avg_p = avg;					// Reset the pointer to the start
			BYTE* pixel = (BYTE*)avg_img->GetPixelAddress(0,avg_img->GetHeight()-1);
			// The pitch is not the same as the image width, therefore we 
			// must check to see if any extra bytes are stored beyond the
			// end of the image row.
			int buffer = abs(avg_img->GetPitch()) % (avg_img->GetBPP()/8);
			for(int i=0; i<avg_img->GetHeight(); i++)
			{
				for(int j=0; j<avg_img->GetWidth(); j++)
				{
					*pixel++ = (BYTE)*avg_p++;
					*pixel++ = (BYTE)*avg_p++;
					*pixel++ = (BYTE)*avg_p++;
				}
				for(int k=0; k<buffer; k++)
				{
					pixel++;
				}
			}
			return avg_img;
		}
		else
		{
			cout << "The images in the set do not have matching dimensions" << endl;
			return NULL;
		}
	}
	else
	{
		cout << "There are no images in the set" << endl;
		return NULL;
	}
}

// This function returns a set of images of the specified 
// dimensions by randomly selecting subwindows from the current
// image set.
ImageSet* ImageSet::CreateRandomWindows(int w, int h, int n, bool rand_scale)
{
	if(_head)
	{
		// Create the new set by generating random windows from
		// the first image.
		Image* rand_img = _head->_image->GetRandomWindow(w,h,rand_scale);
		ImageSet* randset = new ImageSet(rand_img);
		delete rand_img;
		for(int i=0; i<n-1; i++)
		{
			randset->_tail->next = new ImageNode;
			randset->_tail->next->_image = _head->_image->GetRandomWindow(w,h,rand_scale);
			randset->_tail->next->next = NULL;
			randset->_tail->next->previous = randset->_tail;
			randset->_tail = randset->_tail->next;
			randset->_n++;
		}

		// Loop over all the images in this set and generate the 
		// specified number of random windows per image.
		ImageNode* temp = _head->next;
		while(temp)
		{
			for(int i=0; i<n; i++)
			{
				randset->_tail->next = new ImageNode;
				randset->_tail->next->_image = temp->_image->GetRandomWindow(w,h,rand_scale);
				randset->_tail->next->next = NULL;
				randset->_tail->next->previous = randset->_tail;
				randset->_tail = randset->_tail->next;
				randset->_n++;
			}
			temp = temp->next;
		}
		
		// Return the set
		return randset;
	}
	else
	{
		cout << "" << endl;
		return NULL;
	}
}